use data_structure::prelude::*;

struct Spare {
    data: Vec<ThreadedAVLTree<usize, i64>>,
    y: usize,
    x: usize,
}

impl Spare {
    #[inline]
    fn zeros(m: usize, n: usize) -> Self {
        let mut d: Vec<ThreadedAVLTree<usize, i64>> = Vec::with_capacity(m);
        for _ in 0..m {
            d.push(ThreadedAVLTree::new());
        }
        Spare {
            data: d,
            y: m,
            x: n,
        }
    }
    #[inline]
    fn eyes(m: usize) -> Self {
        let mut d: Vec<ThreadedAVLTree<usize, i64>> = Vec::with_capacity(m);
        for i in 0..m {
            let mut tree = ThreadedAVLTree::new();
            tree.insert(i + 1, 1);
            d.push(tree);
        }
        Spare {
            data: d,
            y: m,
            x: m,
        }
    }

    #[inline]
    fn at(&self, m: usize, n: usize) -> i64 {
        assert!(m > 0 && m <= self.y && n > 0 && n <= self.x);
        match self.data[m - 1].query(n) {
            Some(x) => x,
            None => 0,
        }
    }

    #[inline]
    fn set(&mut self, m: usize, n: usize, val: i64) {
        assert!(m > 0 && m <= self.y && n > 0 && n <= self.x);
        if !self.data[m - 1].insert(n, val) {
            self.data[m - 1].update(n, val)
        }
    }

    #[inline]
    fn solve(&self, x: Vec<i64>) -> Vec<i64> {
        assert!(x.len() == self.x);
        let mut ans = Vec::<i64>::with_capacity(self.x);
        for m in 0..self.y {
            let mut accum = 0i64;
            for (n, val) in self.data[m].to_vector().into_iter() {
                accum += x[n - 1] * val;
            }
            ans.push(accum);
        }
        ans
    }
}

#[test]
fn spare_matrix() {
    let mut a = Spare::eyes(10);
    for m in 1..=10 {
        a.set(m, 10, 1);
    }
    let answer = a.solve(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
    assert_eq!(answer, vec![11, 12, 13, 14, 15, 16, 17, 18, 19, 10]);
}
